# import os
# import cv2
# import six
# import lmdb
# import argparse
# from PIL import Image
# import numpy as np
# parser = argparse.ArgumentParser()
# parser.add_argument('--input', type=str, default='DocTamperV1-SCD')  ## path to the dataset dir, which contains data.mdb and lock.mdb
# parser.add_argument('--output', type=str, default='DTV1-SCD')  ## path to the dataset dir, which contains data.mdb and lock.mdb
# parser.add_argument('--total', type=int, default=18000) # index of the image to be visualized
# args = parser.parse_args()
# a =  lmdb.open(args.input,readonly=True,lock=False,readahead=False,meminit=False)
# for index in range(args.total):
# # index = args.i
#     with a.begin(write=False) as txn:
#         img_key = 'image-%09d' % index
#         imgbuf = txn.get(img_key.encode('utf-8'))
#         buf = six.BytesIO()
#         buf.write(imgbuf)
#         buf.seek(0)
#         im = Image.open(buf)
#         img_path = os.path.join(args.output, "img", f"{index}.jpg")
#         im.save(img_path)
#         lbl_key = 'label-%09d' % index
#         lblbuf = txn.get(lbl_key.encode('utf-8'))
#         mask = cv2.imdecode(np.frombuffer(lblbuf,dtype=np.uint8),0)
#         print(index)
#         # print(mask.max())
#         if mask.max()==1:
#             mask=mask*255
#         mask_path = os.path.join(args.output, "mask", f"{index}.png")
#         cv2.imwrite(mask_path,mask)
#
import os
import cv2
import six
import lmdb
import argparse
from PIL import Image
import numpy as np

parser = argparse.ArgumentParser()
parser.add_argument('--input', type=str, default='DocTamperV1-SCD')  ## path to the dataset dir, which contains data.mdb and lock.mdb
parser.add_argument('--output', type=str, default='DTV1-SCD')  ## path to the dataset dir, which contains data.mdb and lock.mdb
parser.add_argument('--total', type=int, default=18000) # index of the image to be visualized
# parser.add_argument('--input', type=str, required=True, help='Path to LMDB dataset')
# parser.add_argument('--output', type=str, required=True, help='Output directory')
# parser.add_argument('--total', type=int, default=18000, help='Number of entries to process')
args = parser.parse_args()

# Create output directories
os.makedirs(os.path.join(args.output, "img"), exist_ok=True)
os.makedirs(os.path.join(args.output, "mask"), exist_ok=True)

env = lmdb.open(args.input,
                readonly=True,
                lock=False,
                readahead=False,
                meminit=False,
                max_readers=256)

for index in range(args.total):
    try:
        with env.begin(write=False) as txn:
            # Process image
            img_key = f'image-{index:09d}'
            imgbuf = txn.get(img_key.encode('utf-8'))
            if not imgbuf:
                print(f"Missing image: {img_key}")
                continue

            with Image.open(six.BytesIO(imgbuf)) as im:
                im.save(os.path.join(args.output, "img", f"{index}.jpg"))

            # Process mask
            lbl_key = f'label-{index:09d}'
            lblbuf = txn.get(lbl_key.encode('utf-8'))
            if not lblbuf:
                print(f"Missing label: {lbl_key}")
                continue

            mask = cv2.imdecode(np.frombuffer(lblbuf, dtype=np.uint8), 0)
            if mask is None:
                raise ValueError("OpenCV failed to decode mask")

            if mask.max() == 1:
                mask = mask * 255

            cv2.imwrite(os.path.join(args.output, "mask", f"{index}.png"), mask)

            print(f"Processed {index}/{args.total}", end='\r')

    except Exception as e:
        print(f"\nError at index {index}: {str(e)}")
        continue

print("\nDone!")